"""
数据集处理模块
包含自定义数据集类和数据加载相关功能
"""

import json
from PIL import Image
from utils.image_download import download_image_from_url
from utils.file_utils import cleanup_temp_files
import torch
from torch.utils.data import Dataset



class ImageFromURLDataset(Dataset):
    """从URL加载图像的自定义数据集，兼容PyTorch DataLoader"""
    def __init__(self, json_path, transform=None):
        with open(json_path, 'r', encoding='utf-8') as f:
            self.data = json.load(f)
        self.transform = transform
        # 创建标签到索引的映射
        self.labels = sorted(list(set(item['image']['point_code'] for item in self.data)))
        self.label_to_idx = {label: idx for idx, label in enumerate(self.labels)}

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        item = self.data[idx]['image']
        image_url = item['url']
        label_str = item['point_code']
        label_idx = self.label_to_idx[label_str]

        try:
            # 使用封装的下载函数，支持HTTP和TOS
            img_fileName = image_url.split('/', 1)[-1] if '/' in image_url else image_url
            download_image_from_url(image_url, max_retries=3, timeout=15)            

        except Exception as e:
            print(f"警告: 无法加载图片 {image_url}. 错误: {e}. 将使用一张占位图.")
            # 创建一个黑色的占位图，防止训练中断
            image = Image.new('RGB', (224, 224), color='black')

        if self.transform:
            try:
                image = Image.open(f'/data/temp/images/{img_fileName}').convert('RGB')
                image = self.transform(image)
            except Exception as e:
                print(f"警告: 图像预处理失败. 错误: {e}. 使用原始图像.")
                # 如果预处理失败，创建一个基础的tensor
                from torchvision import transforms
                image = transforms.ToTensor()(image)
            
        return image, label_idx
        
    def get_num_classes(self):
        return len(self.labels)

    def get_labels(self):
        return self.labels


def get_model_transform(is_training=True):
    """获取标准化的图像预处理transforms"""
    from torchvision import transforms
    
    if is_training:
        return transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])
    else:
        return transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])
